package me.yuxiaoyao.lock.interceptor;

import me.yuxiaoyao.lock.LockExecutor;
import me.yuxiaoyao.lock.exception.DistributedLockException;
import me.yuxiaoyao.lock.exception.LockFailedException;
import me.yuxiaoyao.lock.exception.NullLockKeyException;
import org.springframework.aop.framework.AopProxyUtils;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.*;
import org.springframework.beans.factory.annotation.BeanFactoryAnnotationUtils;
import org.springframework.context.expression.AnnotatedElementKey;
import org.springframework.expression.EvaluationContext;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.util.function.SingletonSupplier;
import org.springframework.util.function.SupplierUtils;

import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.function.Supplier;

/**
 * @author kerryzhang on 2020/12/22
 */


public abstract class DistributedLockAspectSupport extends AbstractDistributedLockInvoker implements BeanFactoryAware, SmartInitializingSingleton {

    private final Map<LockOperationCacheKey, OperationMetadata> metadataCache = new ConcurrentHashMap<>(1024);

    @Nullable
    private DistributedLockOperationSource distributedLockOperationSource;
    @Nullable
    private SingletonSupplier<LockExecutor> lockExecutor;
    private SingletonSupplier<KeyGenerator> keyGenerator;

    private final DistributedLockOperationExpressionEvaluator evaluator = new DistributedLockOperationExpressionEvaluator();

    @Nullable
    private BeanFactory beanFactory;
    private boolean initialized = false;

    public void configure(@Nullable Supplier<LockExecutor> lockExecutor,
                          @Nullable Supplier<KeyGenerator> keyGenerator,
                          @Nullable Supplier<ErrorHandler> errorHandler) {
        this.lockExecutor = SingletonSupplier.ofNullable(lockExecutor);
        this.keyGenerator = new SingletonSupplier<>(keyGenerator, SimpleKeyGenerator::new);
        this.errorHandler = new SingletonSupplier<>(errorHandler, DefaultErrorHandler::new);
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    @Override
    public void afterSingletonsInstantiated() {
        if (getLockExecutor() == null) {
            Assert.state(this.beanFactory != null, "LockExecutor or BeanFactory must be set on cache aspect");
            try {
                // set default
                setLockExecutor(this.beanFactory.getBean(LockExecutor.class));
            } catch (NoUniqueBeanDefinitionException ex) {
                throw new IllegalStateException("No LockExecutor specified, and no unique bean of type " +
                        "LockExecutor found. Mark one as primary or declare a specific LockRegistry to use.");
            } catch (NoSuchBeanDefinitionException ex) {
                throw new IllegalStateException("No LockExecutor specified, and no bean of type LockExecutor found. " +
                        "Register a LockExecutor bean or remove the @EnableDistributedLock annotation from your configuration.");
            }
        }
        this.initialized = true;
    }

    protected Object execute(Invoker invoker, Object target, Method method, Object[] args) throws ThrowableWrapper {
        if (this.initialized) {
            DistributedLockOperationSource source = getDistributedLockOperationSource();
            if (source != null) {
                Class<?> targetClass = getTargetClass(target);
                // 方法上的注解
                Collection<DistributedLockOperation> lockOperations = source.getLockOperations(method, targetClass);
                if (!CollectionUtils.isEmpty(lockOperations)) {
                    List<OperationMetadata> metadataList = getMetadataList(lockOperations, method, targetClass);
                    return execute(invoker, target, targetClass, method, args, metadataList);
                }
            }
        }
        // 直接运行
        return invoker.invoke();
    }


    private Object execute(Invoker invoker, Object target, Class<?> targetClass, Method method, Object[] args, List<OperationMetadata> metadataList) throws ThrowableWrapper {
        List<OperationContext> contexts = getOperationContext(metadataList, target, targetClass, args);

        List<CurrentLockContext> lockContexts = new ArrayList<>(contexts.size());
        OperationContext context = null;
        try {
            for (OperationContext operationContext : contexts) {
                context = operationContext;
                Object key = context.generateKey();
                if (key == null) {
                    throw new NullLockKeyException(context.metadata.operation.toString());
                }
                Lock lock = context.getMetadata().lockExecutor.getLock(key);
                lockContexts.add(new CurrentLockContext(key, lock, context.metadata.operation.isBlock(), context.metadata.operation.getTimeout()));
            }
        } catch (Throwable t) {
            getErrorHandler().handleDistributedLockError(t, context == null ? null : context.getMetadata().operation.toString());
            // 避免不处理异常
            throw new DistributedLockException(context == null ? null : context.getMetadata().operation.toString());
        }
        try {
            // do lock
            try {
                for (CurrentLockContext lockContext : lockContexts) {
                    Lock lock = lockContext.getLock();
                    if (lockContext.isBlock()) {
                        // block Lock
                        if (lockContext.getTimeout() > 0) {
                            if (lock instanceof me.yuxiaoyao.lock.Lock) {
                                // do lock
                                ((me.yuxiaoyao.lock.Lock) lock).lock(lockContext.timeout, TimeUnit.MILLISECONDS);
                            } else {
                                boolean lockStatus = lock.tryLock(lockContext.timeout, TimeUnit.MILLISECONDS);
                                if (!lockStatus) {
                                    lockContext.setLock(false);
                                    throw new LockFailedException(lockContext.getKey().toString());
                                }
                            }
                        } else {
                            //noinspection AlibabaLockShouldWithTryFinally
                            lock.lock();
                        }
                    } else {
                        // unblock Lock
                        boolean lockStatus;
                        if (lockContext.getTimeout() > 0) {
                            lockStatus = lock.tryLock(lockContext.timeout, TimeUnit.MILLISECONDS);
                        } else {
                            lockStatus = lock.tryLock();
                        }
                        if (!lockStatus) {
                            lockContext.setLock(false);
                            throw new LockFailedException(lockContext.getKey().toString());
                        }
                    }
                    System.out.println("lockKey = " + lockContext.getKey().getClass() + ". value = " + lockContext.getKey());

                }
            } catch (InterruptedException | IllegalStateException e) {
                // 有锁失败的情况
                unlockContentList(lockContexts);
                throw new DistributedLockException(e.getMessage(), e);
            } catch (Throwable ex) {
                unlockContentList(lockContexts);
                // 有锁失败的情况
                throw ex;
            }
            // all lock success

            // 执行任务
            return invoker.invoke();
        } finally {
            // unlock
            unlockContentList(lockContexts);
        }
    }

    private void unlockContentList(List<CurrentLockContext> lockContexts) {
        for (int i = lockContexts.size() - 1; i >= 0; i--) {
            if (lockContexts.get(i).isLock) {
                try {
                    lockContexts.get(i).getLock().unlock();
                } catch (Throwable ignored) {
                    // some key unlock failed
                }
            }
        }
    }

    private static class CurrentLockContext {
        private final Object key;
        private final Lock lock;
        private final boolean isBlock;
        private final long timeout;
        private boolean isLock = true;

        private CurrentLockContext(Object key, Lock lock, boolean isBlock, long timeout) {
            this.key = key;
            this.lock = lock;
            this.isBlock = isBlock;
            this.timeout = timeout;
        }

        public Object getKey() {
            return key;
        }

        public Lock getLock() {
            return lock;
        }

        public boolean isBlock() {
            return isBlock;
        }

        public long getTimeout() {
            return timeout;
        }

        public boolean isLock() {
            return isLock;
        }

        public CurrentLockContext setLock(boolean lock) {
            isLock = lock;
            return this;
        }
    }

    private List<OperationContext> getOperationContext(Collection<OperationMetadata> metadataList, Object target, Class<?> targetClass, Object[] args) {
        List<OperationContext> contexts = new ArrayList<>();
        for (OperationMetadata metadata : metadataList) {
            contexts.add(new OperationContext(metadata, args, target));
        }
        return contexts;
    }

    private List<OperationMetadata> getMetadataList(Collection<DistributedLockOperation> operations, Method method, Class<?> targetClass) {
        List<OperationMetadata> metadataList = new ArrayList<>(operations.size());
        for (DistributedLockOperation operation : operations) {
            metadataList.add(getMetadata(operation, method, targetClass));
        }
        return metadataList;
    }


    private OperationMetadata getMetadata(DistributedLockOperation operation, Method method, Class<?> targetClass) {
        // key
        LockOperationCacheKey cacheKey = new LockOperationCacheKey(operation, method, targetClass);
        // value
        OperationMetadata metadata = this.metadataCache.get(cacheKey);
        if (metadata == null) {
            LockExecutor lockExecutor;
            if (StringUtils.hasText(operation.getLockExecutor())) {
                lockExecutor = getBean(operation.getLockExecutor(), LockExecutor.class);
            } else {
                // set default
                lockExecutor = this.getLockExecutor();
            }

            KeyGenerator keyGenerator = null;
            if (StringUtils.hasText(operation.getKeyGenerator())) {
                keyGenerator = getBean(operation.getKeyGenerator(), KeyGenerator.class);
            }

            metadata = new OperationMetadata(operation, method, targetClass, lockExecutor, keyGenerator);
            this.metadataCache.put(cacheKey, metadata);
        }
        return metadata;

    }


    public void setDistributedLockOperationSource(@Nullable DistributedLockOperationSource distributedLockOperationSource) {
        this.distributedLockOperationSource = distributedLockOperationSource;
    }

    @Nullable
    public DistributedLockOperationSource getDistributedLockOperationSource() {
        return distributedLockOperationSource;
    }

    public void setLockExecutor(LockExecutor lockExecutor) {
        this.lockExecutor = SingletonSupplier.ofNullable(lockExecutor);
    }

    @Nullable
    public LockExecutor getLockExecutor() {
        return SupplierUtils.resolve(this.lockExecutor);
    }

    public KeyGenerator getKeyGenerator() {
        return SupplierUtils.resolve(keyGenerator);
    }

    private Class<?> getTargetClass(Object target) {
        return AopProxyUtils.ultimateTargetClass(target);
    }


    protected <T> T getBean(String beanName, Class<T> expectedType) {
        if (this.beanFactory == null) {
            throw new IllegalStateException(
                    "BeanFactory must be set on cache aspect for " + expectedType.getSimpleName() + " retrieval");
        }
        return BeanFactoryAnnotationUtils.qualifiedBeanOfType(this.beanFactory, expectedType, beanName);
    }


    /**
     * wrapper
     */
    protected class OperationContext {
        private final OperationMetadata metadata;
        private final Object[] args;
        private final Object target;

        public OperationContext(OperationMetadata metadata, Object[] args, Object target) {
            this.metadata = metadata;
            this.args = args;
            this.target = target;
        }

        public OperationMetadata getMetadata() {
            return metadata;
        }

        public Object[] getArgs() {
            return args;
        }

        public Object getTarget() {
            return target;
        }


        @Nullable
        protected Object generateKey() {
            if (this.metadata.keyGenerator != null) {
                return this.metadata.keyGenerator.generate(this.target, this.metadata.method, this.args);
            } else {
                EvaluationContext evaluationContext = createEvaluationContext();
                // 生成KEY
                return evaluator.key(this.metadata.operation.getKey(), evaluationContext);
            }
        }

        private EvaluationContext createEvaluationContext() {
            return evaluator.createEvaluationContext(this.metadata.method, this.args,
                    this.target, this.metadata.targetClass, this.metadata.targetMethod, beanFactory);
        }

    }


    private static class OperationMetadata {
        private final DistributedLockOperation operation;
        private final Method method;
        private final Class<?> targetClass;
        private final Method targetMethod;
        private final AnnotatedElementKey methodKey;
        private final LockExecutor lockExecutor;
        private final KeyGenerator keyGenerator;


        private OperationMetadata(DistributedLockOperation operation, Method method, Class<?> targetClass,
                                  LockExecutor lockExecutor, KeyGenerator keyGenerator) {
            this.operation = operation;
            this.method = method;
            this.targetClass = targetClass;
            this.targetMethod = (!Proxy.isProxyClass(targetClass) ? AopUtils.getMostSpecificMethod(method, targetClass) : this.method);
            this.methodKey = new AnnotatedElementKey(this.targetMethod, targetClass);
            this.lockExecutor = lockExecutor;
            this.keyGenerator = keyGenerator;
        }
    }


    private static class LockOperationCacheKey implements Comparable<LockOperationCacheKey> {
        private final DistributedLockOperation operation;

        private final AnnotatedElementKey methodCacheKey;

        private LockOperationCacheKey(DistributedLockOperation operation, Method method, Class<?> targetClass) {
            this.operation = operation;
            this.methodCacheKey = new AnnotatedElementKey(method, targetClass);
        }

        @Override
        public boolean equals(@Nullable Object other) {
            if (this == other) {
                return true;
            }
            if (!(other instanceof LockOperationCacheKey)) {
                return false;
            }
            LockOperationCacheKey otherKey = (LockOperationCacheKey) other;
            return (this.operation.equals(otherKey.operation) &&
                    this.methodCacheKey.equals(otherKey.methodCacheKey));
        }

        @Override
        public int hashCode() {
            return (this.operation.hashCode() * 31 + this.methodCacheKey.hashCode());
        }

        @Override
        public String toString() {
            return this.operation + " on " + this.methodCacheKey;
        }

        @Override
        public int compareTo(LockOperationCacheKey other) {
            int result = this.operation.getName().compareTo(other.operation.getName());
            if (result == 0) {
                result = this.methodCacheKey.compareTo(other.methodCacheKey);
            }
            return result;
        }
    }

}
